import torch
from torchvision.transforms import transforms
from PIL import Image

# 实例化模型
from model import ActionClassificationModel

model = ActionClassificationModel()

# 构建随机输入数据
batch_size = 1
time_steps = 20
input_channels = 3
img_height = 112
img_width = 112
x = torch.randn(batch_size, time_steps, input_channels, img_height, img_width)

# 模型前向传播
with torch.no_grad():
    output = model(x)

# 输出张量的形状
print("Output shape:", output.shape)

# 加载并转换测试图像
transform = transforms.Compose([
    transforms.Resize((img_height, img_width)),
    transforms.ToTensor(),
])
img = Image.open("dataset/val/drink/drink_101/5.jpg")
img = transform(img)

# 构建随机输入数据
x = img.unsqueeze(0).unsqueeze(0).repeat(batch_size, time_steps, 1, 1, 1)

# 模型前向传播
with torch.no_grad():
    output = model(x)

# 输出分类结果
print("Classification result:", output.argmax(dim=-1))
